# -----------------------------------------------------------------------------
# Copyright (c) 2025, Southeast University (China)
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are
# met:
#  - Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
#  - Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#  - Neither the name of the copyright holders nor the names of its
#    contributors may be used to endorse or promote products derived from
#    this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
# A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
# OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
# SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
# LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
# DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
# THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#
# Authors: Jiajie Xu, Yanfang Zhang, Jiaqi Gao, Leyun Tian
# -----------------------------------------------------------------------------


import subprocess
import shutil
import signal
import atexit
from work_var import *

def check_command(command_name):
    # check if the command is available
    return shutil.which(command_name) is not None

def Run_Pt_Script(script):
    """
    Run a PrimeTime (pt_shell) script using subprocess, ensuring child processes
    are terminated when the main process exits or is killed.

    Args:
        script (str): The name of the script to run.
    """
    if not check_command('pt_shell'):
        raise EnvironmentError("Error: 'pt_shell' not found. Please ensure PrimeTime is correctly installed and added to your PATH.")
    
    command = ['pt_shell', '-f', script]
    working_directory = os.path.join(work_dir, 'log/')
    
    # Create the working directory if it doesn't exist
    if not os.path.exists(working_directory):
        os.makedirs(working_directory)
    
    print(f'Running PT script: {script} in {working_directory}')
    
    # Start the subprocess and assign it to a new process group
    process = subprocess.Popen(
        command,
        cwd=working_directory,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        text=True,
        bufsize=1,
        universal_newlines=True,
        preexec_fn=os.setsid  # Create a new process group
    )
    
    # Register cleanup to terminate child processes on exit
    def cleanup():
        if process.poll() is None:  # Check if the process is still running
            print("Terminating PT subprocess...")
            os.killpg(os.getpgid(process.pid), signal.SIGTERM)  # Terminate the process group
            print("PT subprocess terminated.")
            
    # noted, problem may occure here if subprocess is runing under a multi-process pool
    atexit.register(cleanup)  # Ensure cleanup is called on normal program exit

    # Handle signals such as Ctrl+C
    def handle_signal(signum, frame):
        print(f"Received signal {signal.Signals(signum).name}, cleaning up...")
        cleanup()
        exit(0)

    signal.signal(signal.SIGINT, handle_signal)  # Handle Ctrl+C
    signal.signal(signal.SIGTERM, handle_signal)  # Handle kill signals

    try:
        # Read and print the subprocess output and error streams
        for line in process.stdout:
            print(line, end='')  # Print standard output
        for err_line in process.stderr:
            print(err_line, end='')  # Print standard error
    except KeyboardInterrupt:
        print("Keyboard interrupt detected. Terminating PT script...")
        cleanup()
    finally:
        process.wait()  # Wait for the process to complete
    
    # Check the subprocess return code
    print('Return code:', process.returncode)


def Run_Inn_Script(script):
    """
    Run a Innovus (innovus) script using subprocess, ensuring child processes
    are terminated when the main process exits or is killed.

    Args:
        script (str): The name of the script to run.
    """
    if not check_command('innovus'):
        raise EnvironmentError("Error: 'innovus' not found. Please ensure Innovus is correctly installed and added to your PATH.")
    
    command = ['innovus', '-files', script, '-no_gui']
    working_directory = os.path.join(work_dir, 'log/')
    
    # Create the working directory if it doesn't exist
    if not os.path.exists(working_directory):
        os.makedirs(working_directory)
    
    print(f'Running Innovus script: {script} in {working_directory}')
    
    # Start the subprocess and assign it to a new process group
    process = subprocess.Popen(
        command,
        cwd=working_directory,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        text=True,
        bufsize=1,
        universal_newlines=True,
        preexec_fn=os.setsid  # Create a new process group
    )
    
    # Register cleanup to terminate child processes on exit
    def cleanup():
        if process.poll() is None:  # Check if the process is still running
            print("Terminating Innovus subprocess...")
            os.killpg(os.getpgid(process.pid), signal.SIGTERM)  # Terminate the process group
            print("Innovus subprocess terminated.")
    
    # noted, problem may occure here if subprocess is runing under a multi-process pool
    atexit.register(cleanup)  # Ensure cleanup is called on normal program exit

    # Handle signals such as Ctrl+C
    def handle_signal(signum, frame):
        print(f"Received signal {signal.Signals(signum).name}, cleaning up...")
        cleanup()
        exit(0)

    signal.signal(signal.SIGINT, handle_signal)  # Handle Ctrl+C
    signal.signal(signal.SIGTERM, handle_signal)  # Handle kill signals

    try:
        # Read and print the subprocess output and error streams
        for line in process.stdout:
            print(line, end='')  # Print standard output
        for err_line in process.stderr:
            print(err_line, end='')  # Print standard error
    except KeyboardInterrupt:
        print("Keyboard interrupt detected. Terminating Innovus script...")
        cleanup()
    finally:
        process.wait()  # Wait for the process to complete
    
    # Check the subprocess return code
    print('Return code:', process.returncode)

def Gen_Inn_output(design="sasc_top", stage="postRoute_fill_std"):
    """
    Generate Innovus output for given design

    Args:
        design: The name of the design to run
        stage: The design stage for innovus output
    """
    initial_script = os.path.join(inn_scripts_dir, "innovus_output.tcl")
    with open(initial_script, "r") as infile:
        script_content = infile.read()
    
    # update scripts with design name and stage 
    updated_content = script_content.replace("set Design sasc_top", f"set Design {design}")
    updated_content = updated_content.replace("source ../innovus_data/$Design/${Design}_postRoute_fill_std.enc", 
                                            f"source ../innovus_data/$Design/${{Design}}_{stage}.enc")
    updated_content = updated_content.replace("_eco", f"_{stage}")
    tmp_scripts = os.path.join(inn_scripts_dir, f"{design}_innovus_output.tcl")
    with open(tmp_scripts, "w") as outfile:
        outfile.write(updated_content)
    
    # generate innovus output
    Run_Inn_Script(tmp_scripts)
    os.remove(tmp_scripts)

def Gen_PT_output(design="sasc_top", stage="postRoute_fill_std"):
    """
    Generate PrimeTime output for given design

    Args:
        design: The name of the design to run
        stage: The design stage for PT output
    """
    initial_script = os.path.join(pt_scripts_dir, "pt_rpt.tcl")
    with open(initial_script, "r") as infile:
        script_content = infile.read()
    
    # update scripts with design name and stage 
    updated_content = script_content.replace("set top_design sasc_top", f"set top_design {design}")
    updated_content = updated_content.replace("_postRoute_fill_std", f"_{stage}")
    tmp_scripts = os.path.join(pt_scripts_dir, f"{design}_pt_rpt.tcl")
    with open(tmp_scripts, "w") as outfile:
        outfile.write(updated_content)
    
    # generate PT output
    Run_Pt_Script(tmp_scripts)
    os.remove(tmp_scripts)
    
if __name__ == "__main__": 
    # Run_Pt_Script('pt_scripts/pt_rpt.tcl')
    # Gen_Inn_output("mc_top")
    # Gen_PT_output("s27")
    Gen_PT_output("mc_top")
    pass
